from qutip import *
from numpy import *
from numpy.linalg import *
from scipy.optimize import fsolve
from math import sin,cos

II = tensor(qeye(2),qeye(2))

sx = sigmax()
sy = sigmay()
sz = sigmaz()
sxx = tensor(sx,sx)
syy = tensor(sy,sy)
szz = tensor(sz,sz)
Q = Qobj([[1,0,0,1j],[0,1j,1,0],[0,1j,-1,0],[1,0,0,-1j]],dims=sxx.dims,shape=sxx.shape)/sqrt(2)

def rot(si):
    return (1j*si*pi/4).expm()

#decompose operator to its local equilalent operator Ud
def cal_local_equivalent_operator(U):
    dt = det(U.full())
    u = tomagic(U)
    m = u.trans()*u
    print 'detm',det(m.full())
    #calculate gammas
    g1 = m.tr()**2/dt/16.0
    g2 = real((m.tr()**2 - (m**2).tr())/dt/4.0)
    print 'g2,g2:',g1,g2
    g1re = real(g1)
    g1im = imag(g1)
    gammas = fsolve(f,x0=[1,1,1],args=(g1re,g1im,g2),xtol=1e-10).tolist()
    rerange(gammas)
    print gammas
    gammas.sort(reverse=True)
    print gammas
    print f(gammas,g1re,g1im,g2)
    diag = [exp(1j*(gammas[0]-gammas[1]+gammas[2])),\
            exp(1j*(gammas[0]+gammas[1]-gammas[2])),\
            exp(-1j*(gammas[0]+gammas[1]+gammas[2])),\
            exp(1j*(-gammas[0]+gammas[1]+gammas[2]))]

    #calculate kl,kr
    evals,estates = m.eigenstates()
    Or = []
    for estate in estates:
        Or.append(estate.trans().full().tolist()[0])
    print 'det Or',det(Or)
    Or = Qobj(Or,dims = u.dims,shape=u.shape)
#     print Or
    kr = tocomp(Or)
#     print kr
    Ud = (0.5j*(gammas[0]*sxx+gammas[1]*syy+gammas[2]*szz)).expm()
    kl = U *kr.dag()*Ud.dag()
    return Ud,kl,kr,gammas
    
    
 
#gamma is periodic with pi, and symmetric by pi/2
def rerange(gammas):
    for i in range(len(gammas)):
        n = int(floor(gammas[i]/pi))
        gammas[i] = gammas[i]-n*pi
        while gammas[i] >= pi:
            gammas[i] = gammas[i]-pi
        while gammas[i] < 0 :
            gammas[i] = gammas[i]+pi
        if gammas[i] > pi/2:
            gammas[i]=pi - gammas[i] 
            
"""converto U to magic basis"""
def tomagic(U):
    return Q.dag()*U*Q
"""convert U to computational basis"""
def tocomp(U):
    return Q*U*Q.dag()

"""function for solve equation to calculate gammas"""
def f(x,g1re,g1im,g2):
    gamma1,gamma2,gamma3 = x
    return [G1_real_func(gamma1,gamma2,gamma3,g1re),\
            G1_imag_func(gamma1,gamma2,gamma3,g1im),\
            G2_func(gamma1,gamma2,gamma3,g2)]

def G1_real_func(gamma1,gamma2,gamma3,g1re):
    return cos(gamma1)**2*cos(gamma2)**2*cos(gamma3)**2-sin(gamma1)**2*sin(gamma2)**2*sin(gamma3)**2 -g1re

def G1_imag_func(gamma1,gamma2,gamma3,g1im):
    return sin(gamma1*2)*sin(gamma2*2)*sin(gamma3*2) -4*g1im

def G2_func(gamma1,gamma2,gamma3,g2):
    return 4*(cos(gamma1)**2*cos(gamma2)**2*cos(gamma3)**2-sin(gamma1)**2*sin(gamma2)**2*sin(gamma3)**2)\
        -cos(gamma1*2)*cos(gamma2*2)*cos(gamma3*2)-g2

kA=tensor(rot(sy),rot(-sy))
kB=tensor(rot(sx),rot(sx))
kC=tensor(rot(sz),rot(sz))
kD=tensor(rot(sz),rot(-sz))

"""merge kl,kr,kA... opeartors to k0,k1,k2,k3"""
def merge_operatiors(kl,kr,k):
    k0 = k.dag()*kr
    k1 = k.dag()*kA.dag()*kB.dag()*k
    k2 = k.dag()*kD.dag()*kB.dag()*kC.dag()*kB*kA*k
    k3 = kl*kC*kB*kD*k
    return k0,k1,k2,k3

"""verification local equivalence decomposition"""
def verify_led(targetU,gammas,H,k0,k1,k2,k3):
    print 'level1 verification local equivalence decomposition...'
    t1 = gammas[0]
    t2 = - gammas[1]
    t3 = - gammas[2]
    U=k3*(1j*t3*H).expm()*k2*(1j*t2*H).expm()*k1*(1j*t1*H).expm()*k0
    
    fidelity = (U.dag()*targetU).tr()/4.0
    print 'verify fidelity:',fidelity
    
def decompose2products(U):
    evals,estates = U.eigenstates()
    logevals = log(evals)
    
    logmatrix = 0
    dim = U.shape[0]
    for i in range(len(logevals)):
        logmatrix=logmatrix+logevals[i]*basis(dim,i)*basis(dim,i).dag()
        
    logmatrix.dims = U.dims
    logmatrix.shape = U.shape
    
    sz = sigmaz()
    II = qeye(2)
    k1 = (tensor(sz,II)*logmatrix).tr()/4
    k2 = (tensor(II,sz)*logmatrix).tr()/4 
    k12 = (tensor(sz,sz)*logmatrix).tr()/4 
    print 'k12:',k12
    print k1,k2
    U1=(k1*tensor(sz,II)).expm().transform(estates,inverse=True)
    U2=(k2*tensor(II,sz)).expm().transform(estates,inverse=True)
    return U1.ptrace([0])/2,U2.ptrace([1])/2

def test_decompose2products():
    sx = sigmax()
    sy = sigmay()
    II = qeye(2)
    U1=(-1j*sx*pi/9).expm()
    U2=(-1j*sx*pi/9).expm()  
    u1,u2=decompose2products(tensor(U1,U2))
    fide1 = (U1.dag()*u1).tr()/2
    fide2 = (U2.dag()*u2).tr()/2
    print "test decomposing fidelity:",fide1,fide2
    
from rotating_frame import *

if __name__=='__main__':
    U = swap()
    Ud,kl,kr,gammas = cal_local_equivalent_operator(U)
    H = (szz)/2.0
    k = tensor(rot(-sy),rot(-sy))
    k0,k1,k2,k3=merge_operatiors(kl,kr,k)
    verify_led(U,gammas,H,k0,k1,k2,k3)
    k1.tidyup(1e-5)
    print 'k1:'
    print_matrix(k1.full())
    print 'k2:'
    print_matrix(k2.full())
    
    